%-------------------------------------------------------------------------------
% $Id: wav2d_finite_diff.m,v 1.5 2011/06/29 01:14:20 paul Exp $
% $Date: 2011/06/29 01:14:20 $
% $Author: paul $
%
%% Explicit Finite Difference Solution to 2D wave equation on a membrane
%
% du{tt} = (c^2)*[du{xx} + du{yy}]  
% 
% Domain: -1 =< x,y =< +1, 0 =< t =< 2,
% Assumptions: c = 1.0
% Initial Conditions: 
%                     r^2 = x^2 + y^2
%                     u(x,y,0) = f(x,y) = f(r)
%                     f(r) = exp(-20*r^2)
%                     du(x,y,0){t} = g(x,y) = 0
% Boundary Conditions:u(0,y,t) = u(L,y,t) = u(x,0,t) = u(x,L,t) = 0
%-------------------------------------------------------------------------------
function wav2d_finite_diff(arg_M, arg_N, arg_boundary, arg_graphics)
% Explicit Finite Difference Solution to 1D wave equation

if nargin < 4
  M = 256;                       % number of samples on x and y axis
  N = 800;                       % number of samples in time interval T
  boundary_type = 0;               
  graphic_type = 0;
  msg = sprintf('INFO:\tUsing the default arguments M=%d N=%d', M,N);
  disp(msg);
else           
  M = arg_M;
  N = arg_N;
  boundary_type = arg_boundary;
  graphic_type = arg_graphics;
  msg = sprintf('INFO:\tUser arguments M=%d N=%d"', M,N);
  disp(msg);
end

T = 2;                         % time interval (0,T)
c = 1.0;                      
length = 1.0;                  % length of square membrane edges
                               
h = length / M;                % dx
k = T / N;                     % dt
lamda = c * k/h;               % lamda =< 1 for stability

msg = sprintf('INFO:\tParameters for 2D Finite Difference Solution:');
disp(msg); 
msg = sprintf('INFO:\tT=%f c=%f h=%f k=%f lamda=%f',T,c,h,k,lamda);
disp(msg); 
if (graphic_type == 0)
  disp('INFO:   Graphic display: imagesc with gray scale');
else
  disp('INFO:   Graphic display: mesh or surface');
end
if (boundary_type == 0)
  disp('INFO:   Boundary conditions: force edges of square perimeter to 0');
else
  disp('INFO:   Boundary conditions: force circular boundary to 0 for testing');
end
prompt = sprintf('INPUT:\tPress return to continue');
response=input(prompt);

if (lamda > (sqrt(2)/2))
  warning = sprintf('WARNING:lamda = %f. PDE solution unstable if lamda > sqrt(2)/2', lamda);
  disp(warning);
  hint = sprintf('INFO:\tAdjust parameters M, N, c and T for convergence');
  disp(hint);
  prompt = sprintf('INPUT:\tPress return to continue');
  response=input(prompt);
end

%-------------------------------------------------------------------------------
% Setup AVI file and parameters
aviobj = avifile('wav2d_finite_diffs.avi');
aviobj.compression = 'None';   % no other option for Unix
aviobj.fps = 5;                % frames per second
aviobj.quality = 10;           % low quality is OK

% Fancy surface plot setup
camlight left;                 % light source for Phong shading
lighting phong;

%-------------------------------------------------------------------------------
% Animation loop
fig = figure;

lamda_sq = lamda ^ 2;              % precompute values outside of loop!
factor = 1 - 2*lamda_sq;           
factor2 = 2*factor;

mask = make_mask(M,boundary_type); % mask sets the boundary conditions
                                   % type:0 = square (assignment),
                                   % type:1 = circular (test)

z = zeros(M+1);                    % initialize frame and delayed versions
z1 = z;
z2 = z;

for j=1:N
  if (j==1)                        % set up initial conditions of vibration
    x = linspace(-1,1,M+1);        % construct initial disturbance function z
    m=[ ];
    for k=1:M+1
      m = [m;x];
    end
    r2 = (m.*m)+(m'.*m');
    z = exp(-20*r2) .* mask;       % initial disturbance pulse
    g = zeros(M+1);                % zero initial velocity
  
  elseif (j==2)                    % first time iteration    
    % accumulate shifted frames together
    z = zeros(M+1);    
    z(2:M+1,1:M+1) = z1(1:M  ,1:M+1);                   % shift i+1
    z(1:M  ,1:M+1) = z(1:M  ,1:M+1) + z1(2:M+1,1:M+1);  % shift i-1        
    z(1:M+1,1:M)   = z(1:M+1,1:M)   + z1(1:M+1,2:M+1);  % shift j+1
    z(1:M+1,2:M+1) = z(1:M+1,2:M+1) + z1(1:M+1,1:M);    % shift j-1    
    
    z = (factor*z1 + (lamda_sq/2)*z.*mask + k*g) .* mask;    
  
  else                             % remaining time iterations
    z = zeros(M+1);    
    z(2:M+1,1:M+1) = z1(1:M  ,1:M+1);                   % shift i+1
    z(1:M  ,1:M+1) = z(1:M  ,1:M+1) + z1(2:M+1,1:M+1);  % shift i-1    
    z(1:M+1,1:M)   = z(1:M+1,1:M)   + z1(1:M+1,2:M+1);  % shift j+1
    z(1:M+1,2:M+1) = z(1:M+1,2:M+1) + z1(1:M+1,1:M);    % shift j-1    

    z = (factor2*z1 + lamda_sq*z.*mask - z2) .* mask;        
  end
  
  % update delayed versions of membrane displacement state (amplitude) z 
  z2 = z1;
  z1 = z;  

  graphics(x,x,z,graphic_type);  % different options for displaying the image

  title('2D Wave Equation Finite Differences Solution');
  frame = getframe(fig); 
  aviobj = addframe(aviobj,frame);
  
  debug=0;
  if debug == 1
     prompt = sprintf('INPUT:\tPress return to continue');
     response=input(prompt);
  end
end 

%-------------------------------------------------------------------------------
% Clean up rendering frame and AVI file

close(fig);
aviobj = close(aviobj);

%-------------------------------------------------------------------------------
% Utility functions
%-------------------------------------------------------------------------------

function mask = make_mask(M, type)
  if type == 0
    % construct square boundary mask with the edges set to zero
    mask = zeros(M+1);                 
    mask(2:M,2:M) = ones(M-1);         
  else
    % experiment with a circular boundary mask to see change in
    % behaviour of the wave reflection at the edge
    
    mask = zeros(M+1); 
    x = linspace(-1,1,M+1); 
    for i=1:M+1
      for j=1:M+1
        if x(i)^2+x(j)^2 < 1
          mask(i,j) = 1.0;
        end
      end
    end
  end
  
function graphics(x,y,z,type)
  if type == 0
    imagesc(z); colormap(gray);
  else
    if type == 1    
      mesh(x,x,z,'direct'); 
;      surf(x,x,z,'FaceColor','interp','EdgeColor','none','FaceLighting','phong');            
    end
    axis([-1 1 -1 1 -1 1]);
    daspect([1 1 1]);
    view(30,15); % view at zero elevation debugging     
    colormap hot;
    colorbar;
    grid off;
  end




